import sys
import io
from pathlib import Path
from tqdm import tqdm

def save_manifest(path_list, path_save):
    with io.FileIO(str(path_save), "w") as f:
        for item in path_list:
            f.write((item + "\n").encode('utf-8'))

if __name__ == "__main__":

    path_data = Path("/mnt/data/DATA/LRW/processing")

    path_manifest_train = Path("data/LRW/manifest_train_mfcc_deep_62000_n2.csv")
    path_manifest_valid = Path("data/LRW/manifest_valid_mfcc_deep_62000_n2.csv")
    path_manifest_test = Path("data/LRW/manifest_test_mfcc_deep_62000_n2.csv")
    path_manifest_train.parent.mkdir(parents=True, exist_ok=True)

    num_train_file = 62500
    num_train_count = 0

    # for each word
    v_path_test = ["2"]
    v_path_train = ["2"]
    v_path_valid = ["2"]

    save_flag = False
    v_ids = sorted(path_data.glob("*"))
    print(f"Have {len(v_ids)} dir files, \n {v_ids[:3]}")

    for i, ipath_id in enumerate(v_ids):
        id_name = ipath_id.stem

        dir_video = path_data.joinpath(id_name)

        print(f"\nid= {id_name}")
        v_dir_videos = sorted(dir_video.glob("*/*/video/video_orignal.mp4"))
        print(f"\nid= {id_name}, Have {len(v_dir_videos)} video files, \n {v_dir_videos[:3]}")

        # for each video
        for ipath_video in v_dir_videos:
            
            ipath_data = ipath_video.parent.parent
            #print(f"ipath_data = {ipath_data}")

            ipath_frames = ipath_data.joinpath("frames")
            ipath_landmk = ipath_data.joinpath("landmk")
            #ipath_landmk_bbox = ipath_data.joinpath("landmk_bbox")

            #ipath_deep = ipath_data.joinpath("audio/deep_feat.npy")
            ipath_mfcc = ipath_data.joinpath("audio/mfcc_feat.npy")
            #ipath_coeff = ipath_data.joinpath("video/coeff_rgbface.npy")  
            ipath_coeff = ipath_data.joinpath("video/coeff_reconstruction.npy")
            #ipath_landk = ipath_data.joinpath("video/landmark_rgbface_crop.npy")
            #ipath_landk2= ipath_data.joinpath("video/landmark_rgbface_proj.npy")

            if ipath_coeff.exists() and ipath_mfcc.exists():
                item =  f"{str(ipath_mfcc)}\t{str(ipath_coeff)}"
                num_train_count += 1

                if "train" in str(ipath_data):
                    v_path_train.append(item)
                elif "test" in str(ipath_data):
                    v_path_test.append(item)
                elif "val" in str(ipath_data):
                    v_path_valid.append(item)
                else:
                    raise(f"train/test/val should in the path {ipath_data}")

            if num_train_count % 1000 == 0: 
                print(f"num_train_count = {num_train_count}")


        print(f"num_train_count = {num_train_count}")

        # if vidoes file is ok
        if num_train_count > num_train_file:
            save_flag = True
            save_manifest(v_path_train, path_manifest_train)
            save_manifest(v_path_test, path_manifest_test)
            save_manifest(v_path_valid, path_manifest_valid)

            print(f"create manifest {num_train_count} videos, {len(v_path_train)} train frames -> {path_manifest_train}")
            break

    # if not have 
    if not save_flag:
        save_manifest(v_path_train, path_manifest_train)
        save_manifest(v_path_test, path_manifest_test)
        save_manifest(v_path_valid, path_manifest_valid)
        print(f"create manifest {num_train_count} videos, {len(v_path_train)} train frames -> {path_manifest_train}")
